﻿#ifndef YOLOV8_H
#define YOLOV8_H
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
#include <exception>
#include <vector>
#include <chrono>


#ifndef INPUT_SIZE
#define INPUT_SEG_H 640
#define INPUT_SEG_W 640
#endif // !input_size


namespace seg
{
    struct PreParam
    {
        float ratio = 1.0f;
        float dw = 0.0f;
        float dh = 0.0f;
        float height = 0;
        float width = 0;
    };

    struct OutputSeg {
        int label;
        float prob;
        cv::Rect rect;
        cv::Mat boxMask;
    };
}


typedef std::chrono::high_resolution_clock Time;
typedef std::chrono::duration<double, std::ratio<1, 1000>> ms;
typedef std::chrono::duration<float> fsec;






class YOLOV8
{
    const std::vector<std::string> CLASS_NAMES = {
        "person", "bicycle", "car", "motorcycle", "airplane", "bus",
        "train", "truck", "boat", "traffic light", "fire hydrant",
        "stop sign", "parking meter", "bench", "bird", "cat",
        "dog", "horse", "sheep", "cow", "elephant",
        "bear", "zebra", "giraffe", "backpack", "umbrella",
        "handbag", "tie", "suitcase", "frisbee", "skis",
        "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
        "skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
        "cup", "fork", "knife", "spoon", "bowl",
        "banana", "apple", "sandwich", "orange", "broccoli",
        "carrot", "hot dog", "pizza", "donut", "cake",
        "chair", "couch", "potted plant", "bed", "dining table",
        "toilet", "tv", "laptop", "mouse", "remote",
        "keyboard", "cell phone", "microwave", "oven",
        "toaster", "sink", "refrigerator", "book", "clock", "vase",
        "scissors", "teddy bear", "hair drier", "toothbrush"
    };


    const std::vector<std::vector<unsigned int>> COLORS = {
        { 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
        { 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
        { 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
        { 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
        { 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
        { 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
        { 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
        { 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
        { 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
        { 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
        { 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
        { 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
        { 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
        { 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
        { 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
        { 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
        { 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
        { 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
        { 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
        { 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
        { 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
        { 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
        { 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
        { 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
        { 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
        { 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
        { 80, 183, 189 }, { 128, 128, 0 }
    };


    const std::vector<std::vector<unsigned int>> MASK_COLORS = {
        { 255, 56, 56 }, { 255, 157, 151 }, { 255, 112, 31 },
        { 255, 178, 29 }, { 207, 210, 49 }, { 72, 249, 10 },
        { 146, 204, 23 }, { 61, 219, 134 }, { 26, 147, 52 },
        { 0, 212, 187 }, { 44, 153, 168 }, { 0, 194, 255 },
        { 52, 69, 147 }, { 100, 115, 255 }, { 0, 24, 236 },
        { 132, 56, 255 }, { 82, 0, 133 }, { 203, 56, 255 },
        { 255, 149, 200 }, { 255, 55, 199 }
    };

public:
    YOLOV8();
    ~YOLOV8();
    bool Initialization(const char* onnxfile);
    void detect(cv::Mat& frame);
    const char* NetPath[2] =
    {
        "./yolov8n-seg.transd_simple.onnx",
        "./yolov8s-seg.transd_simple.onnx",
    };

    void draw_objects(cv::Mat& image);


private:
    Ort::Env* env = nullptr;
    Ort::SessionOptions* sessionOptions = nullptr;
    Ort::Session* session = nullptr;

    static constexpr size_t inputTensorSize = 1 * 3 * 640 * 640;
    std::vector<int64_t> inputTensorShape{ 1, 3, 640, 640};
    float* blob = nullptr;

    std::vector<const char*> inputNames = {"images"};
    std::vector<const char*> outputNames = {"output0","output1"};

    int preWidth = 0;
    int preHeight = 0;

    static constexpr int num_anchors = 8400;
    static constexpr int fea_channels = 116;

    static constexpr int segWidth = 160;
    static constexpr int segHeight = 160;
    static constexpr int seg_channels = 32;

    static constexpr int INPUT_SIZE = 1 * 3 * INPUT_SEG_H * INPUT_SEG_W;
    static constexpr int OUTPUT0_SIZE = 1 * fea_channels * num_anchors;
    static constexpr int OUTPUT1_SIZE = 1 * seg_channels * segWidth * segHeight;
    static constexpr int num_classes = fea_channels - seg_channels - 4;

    const cv::Size dst_size = cv::Size(INPUT_SEG_W, INPUT_SEG_H);
    const cv::Scalar color = cv::Scalar(114, 114, 114);
    const int topk = 100;
    const float _classThreshold = 0.320f;
    const float _nmsThreshold = 0.350f;
    const float _maskThreshold = 0.480f;
    cv::Mat M;
    seg::PreParam pparam;

    std::vector<std::vector<float>> mask_confs;
    std::vector<seg::OutputSeg> outputs;
    std::vector<int> labels;
    std::vector<float> scores;
    std::vector<cv::Rect> bboxes;
    std::vector<int> indices;
};


#endif // YOLOV5_H
